import cv2
import numpy as np

from ultralytics import YOLO

from collections import defaultdict


#加载模型
model = YOLO("yolo11s-pose.pt")  # load an official model
model = YOLO("./runs/pose/v11s-pose_epoch100/weights/best.pt")  # load a custom model

#配置视频相关的
video_path = "../materials/dog25fps.mp4"
# video_path = "../petvideo/dog25fps.mp4"
result_path = "../detectvideo/track_dog1.mp4"

#记录所有的id点信息
track_history = defaultdict(lambda :[])

if __name__ == '__main__':
    #抓取视频
    capture = cv2.VideoCapture(video_path)
    if not capture.isOpened():
        print("Error opening Video file.")
        exit()
    #视频参数
    fps=capture.get(cv2.CAP_PROP_FPS)
    frame_width = capture.get(cv2.CAP_PROP_FRAME_WIDTH)
    frame_height = capture.get(cv2.CAP_PROP_FRAME_HEIGHT)

    #窗口等待时间
    wtime = (int)((1/fps)*1000)

    videoWriter =None

    while True:
        successCode,frame = capture.read()
        if not successCode:
            print("视频读取结束")
            break
        #预测单张图片
        results = model.track(frame,persist=True)
        a_frame = results[0].plot()
        a_frame = np.copy(a_frame)

        #所有id的序列号信息
        boxes = results[0].boxes.xywh.cpu()
        if results[0].boxes.id is None:
            track_ids=[]
        else:
            track_ids = results[0].boxes.id.int().cpu().tolist()

        #迭代每一个预测到的对象
        for box,track_id in zip(boxes,track_ids):

            #获得框的大小
            x,y,w,h = box
            track = track_history[track_id]
            #添加坐标
            track.append((float(x),float(y)))
            if len(track)>75:
                track.pop(0)

            #创建点实例
            points = np.hstack(track).astype(np.int32).reshape(-1,1,2)
            #画到预测结果上
            cv2.polylines(a_frame,[points],isClosed=False,color=(0,0,255),thickness=3)

            if videoWriter is None:
                # fourcc = cv2.VideoWriter_fourcc()
                fourcc = cv2.VideoWriter_fourcc(*"mp4v")
                videoWriter = cv2.VideoWriter(result_path, fourcc, fps, (int(frame_width), int(frame_height)))
            videoWriter.write(a_frame)
            cv2.imshow("dog_track",a_frame)
            cv2.waitKey(wtime)
    capture.release()  # 释放视频资源
    #保存视频
    videoWriter.release()
    cv2.destroyAllWindows()


